package com.flk.flkelasticsearch.es;

import org.elasticsearch.action.ActionFuture;
import org.elasticsearch.action.get.GetResponse;
import org.elasticsearch.action.index.IndexRequestBuilder;
import org.elasticsearch.action.search.SearchRequestBuilder;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.client.Client;
import org.elasticsearch.client.Requests;
import org.elasticsearch.search.aggregations.AbstractAggregationBuilder;
import org.elasticsearch.search.fetch.subphase.highlight.HighlightBuilder;
import org.elasticsearch.search.sort.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.data.domain.Sort;
import org.springframework.data.elasticsearch.ElasticsearchException;
import org.springframework.data.elasticsearch.core.DefaultResultMapper;
import org.springframework.data.elasticsearch.core.ElasticsearchTemplate;
import org.springframework.data.elasticsearch.core.ResultsExtractor;
import org.springframework.data.elasticsearch.core.SearchResultMapper;
import org.springframework.data.elasticsearch.core.aggregation.AggregatedPage;
import org.springframework.data.elasticsearch.core.convert.ElasticsearchConverter;
import org.springframework.data.elasticsearch.core.facet.FacetRequest;
import org.springframework.data.elasticsearch.core.mapping.ElasticsearchPersistentEntity;
import org.springframework.data.elasticsearch.core.mapping.ElasticsearchPersistentProperty;
import org.springframework.data.elasticsearch.core.query.*;
import org.springframework.util.Assert;
import org.springframework.util.StringUtils;

import java.io.IOException;
import java.util.List;

import static org.elasticsearch.index.VersionType.EXTERNAL;
import static org.elasticsearch.index.VersionType.EXTERNAL_GTE;
import static org.springframework.util.CollectionUtils.isEmpty;

/**
 * @author xxwy
 * on 2018/11/12 0012
 */

public class MyElasticSearchTemplate extends ElasticsearchTemplate {


    private Client client;
    private static final String FIELD_SCORE = "_score";
    private static final Logger QUERY_LOGGER = LoggerFactory.getLogger("org.springframework.data.elasticsearch.core.QUERY");
    private String searchTimeout;


    public MyElasticSearchTemplate(Client client){
        super(client);
        this.client = client;
    }

    public MyElasticSearchTemplate(Client client, ElasticsearchConverter elasticsearchConverter) {
        super(client, elasticsearchConverter, new DefaultResultMapper(elasticsearchConverter.getMappingContext()));
        this.client = client;
    }

    @Override
    public <T> AggregatedPage<T> queryForPage(SearchQuery query, Class<T> clazz, SearchResultMapper mapper) {
        SearchResponse response = dosearch(prepareSearch(query, clazz), query);
        return mapper.mapResults(response, clazz, query.getPageable());
    }


    public <T> T queryForObject(GetQuery query, String indexName,Class<T> clazz) {
        GetResponse response = client
                .prepareGet(indexName, EsInfomation.DOC, query.getId()).execute()
                .actionGet();
        T entity = getResultsMapper().mapResult(response, clazz);
        return entity;
    }

    @Override
    public String index(IndexQuery query) {
        String documentId = prepareIndex(query).execute().actionGet().getId();
        // We should call this because we are not going through a mapper.
        if (query.getObject() != null) {
            setPersistentEntityId(query.getObject(), documentId);
        }
        return documentId;
    }

    public IndexRequestBuilder prepareIndex(IndexQuery query) {
        try {
            String indexName = StringUtils.isEmpty(query.getIndexName())
                    ? retrieveIndexNameFromPersistentEntity(query.getObject().getClass())[0]
                    : query.getIndexName();
            String type = StringUtils.isEmpty(query.getType()) ? retrieveTypeFromPersistentEntity(query.getObject().getClass())[0]
                    : query.getType();

            IndexRequestBuilder indexRequestBuilder = null;

            if (query.getObject() != null) {
                String id = StringUtils.isEmpty(query.getId()) ? getPersistentEntityId(query.getObject()) : query.getId();
                // If we have a query id and a document id, do not ask ES to generate one.
                if (id != null) {
                    indexRequestBuilder = client.prepareIndex(indexName, type, id);
                } else {
                    indexRequestBuilder = client.prepareIndex(indexName, type);
                }
                indexRequestBuilder.setSource(getResultsMapper().getEntityMapper().mapToString(query.getObject()),
                        Requests.INDEX_CONTENT_TYPE);
            } else if (query.getSource() != null) {
                indexRequestBuilder = client.prepareIndex(indexName, type, query.getId()).setSource(query.getSource(),
                        Requests.INDEX_CONTENT_TYPE);
            } else {
                throw new ElasticsearchException(
                        "object or source is null, failed to index the document [id: " + query.getId() + "]");
            }
            if (query.getVersion() != null) {
                indexRequestBuilder.setVersion(query.getVersion());
                //todo:修改过的代码
                indexRequestBuilder.setVersionType(EXTERNAL_GTE);
            }

            if (query.getParentId() != null) {
                indexRequestBuilder.setParent(query.getParentId());
            }

            return indexRequestBuilder;
        } catch (IOException e) {
            throw new ElasticsearchException("failed to index the document [id: " + query.getId() + "]", e);
        }
    }

    private void setPersistentEntityId(Object entity, String id) {

        ElasticsearchPersistentEntity<?> persistentEntity = getPersistentEntityFor(entity.getClass());
        ElasticsearchPersistentProperty idProperty = persistentEntity.getIdProperty();

        // Only deal with text because ES generated Ids are strings !

        if (idProperty != null && idProperty.getType().isAssignableFrom(String.class)) {
            persistentEntity.getPropertyAccessor(entity).setProperty(idProperty, id);
        }
    }

    private String getPersistentEntityId(Object entity) {

        ElasticsearchPersistentEntity<?> persistentEntity = getPersistentEntityFor(entity.getClass());
        Object identifier = persistentEntity.getIdentifierAccessor(entity).getIdentifier();

        if (identifier != null) {
            return identifier.toString();
        }

        return null;
    }

        @Override
    public <T> T query(SearchQuery query, ResultsExtractor<T> resultsExtractor) {
        SearchResponse response = dosearch(prepareSearch(query), query);
        return resultsExtractor.extract(response);
    }

    private SearchResponse dosearch(SearchRequestBuilder searchRequest, SearchQuery searchQuery) {
        if (searchQuery.getFilter() != null) {
            searchRequest.setPostFilter(searchQuery.getFilter());
        }

        if (!isEmpty(searchQuery.getElasticsearchSorts())) {
            for (SortBuilder sort : searchQuery.getElasticsearchSorts()) {
                searchRequest.addSort(sort);
            }
        }

        if (!searchQuery.getScriptFields().isEmpty()) {
            // _source should be return all the time
            // searchRequest.addStoredField("_source");
            for (ScriptField scriptedField : searchQuery.getScriptFields()) {
                searchRequest.addScriptField(scriptedField.fieldName(), scriptedField.script());
            }
        }

        if (searchQuery.getHighlightFields() != null || searchQuery.getHighlightBuilder() != null) {
            HighlightBuilder highlightBuilder = searchQuery.getHighlightBuilder();
            if (highlightBuilder == null) {
                highlightBuilder = new HighlightBuilder();
            }
            for (HighlightBuilder.Field highlightField : searchQuery.getHighlightFields()) {
                highlightBuilder.field(highlightField);
            }
            searchRequest.highlighter(highlightBuilder);
        }

        if (!isEmpty(searchQuery.getIndicesBoost())) {
            for (IndexBoost indexBoost : searchQuery.getIndicesBoost()) {
                searchRequest.addIndexBoost(indexBoost.getIndexName(), indexBoost.getBoost());
            }
        }

        if (!isEmpty(searchQuery.getAggregations())) {
            for (AbstractAggregationBuilder aggregationBuilder : searchQuery.getAggregations()) {
                searchRequest.addAggregation(aggregationBuilder);
            }
        }

        if (!isEmpty(searchQuery.getFacets())) {
            for (FacetRequest aggregatedFacet : searchQuery.getFacets()) {
                searchRequest.addAggregation(aggregatedFacet.getFacet());
            }
        }
        return getSearchResponse(searchRequest.setQuery(searchQuery.getQuery()));
    }

    private SearchResponse getSearchResponse(SearchRequestBuilder requestBuilder) {

        if (QUERY_LOGGER.isDebugEnabled()) {
            QUERY_LOGGER.debug(requestBuilder.toString());
        }

        return getSearchResponse(requestBuilder.execute());
    }

    private SearchResponse getSearchResponse(ActionFuture<SearchResponse> response) {
        return searchTimeout == null ? response.actionGet() : response.actionGet(searchTimeout);
    }



    private  <T> SearchRequestBuilder prepareSearch(Query query, Class<T> clazz) {
        if (query.getIndices().isEmpty()) {
            query.addIndices(retrieveIndexNameFromPersistentEntity(clazz));
        }
        if (query.getTypes().isEmpty()) {
            query.addTypes(retrieveTypeFromPersistentEntity(clazz));
        }
        return prepareSearch(query);
    }




    private String[] retrieveIndexNameFromPersistentEntity(Class clazz) {
        if (clazz != null) {
            return new String[] { getPersistentEntityFor(clazz).getIndexName() };
        }
        return null;
    }

    private String[] retrieveTypeFromPersistentEntity(Class clazz) {
        if (clazz != null) {
            return new String[] { getPersistentEntityFor(clazz).getIndexType() };
        }
        return null;
    }


    private SearchRequestBuilder prepareSearch(Query query) {
        Assert.notNull(query.getIndices(), "No index defined for Query");
        Assert.notNull(query.getTypes(), "No type defined for Query");

        int startRecord = 0;
        SearchRequestBuilder searchRequestBuilder = client.prepareSearch(toArray(query.getIndices()))
                .setSearchType(query.getSearchType())
                .setTypes(toArray(query.getTypes()))
                .setVersion(true)
                .setTrackScores(query.getTrackScores());

        if (query.getSourceFilter() != null) {
            SourceFilter sourceFilter = query.getSourceFilter();
            searchRequestBuilder.setFetchSource(sourceFilter.getIncludes(), sourceFilter.getExcludes());
        }

        if (query.getPageable().isPaged()) {
            startRecord = query.getPageable().getPageNumber() * query.getPageable().getPageSize();
            searchRequestBuilder.setSize(query.getPageable().getPageSize());
        }
        searchRequestBuilder.setFrom(startRecord);

        if (!query.getFields().isEmpty()) {
            //todo:修改过的代码
            searchRequestBuilder.storedFields(toArray(query.getFields()));
        }

        if (query.getIndicesOptions() != null) {
            searchRequestBuilder.setIndicesOptions(query.getIndicesOptions());
        }

        if (query.getSort() != null) {
            for (Sort.Order order : query.getSort()) {
                SortOrder sortOrder = order.getDirection().isDescending() ? SortOrder.DESC : SortOrder.ASC;

                if (FIELD_SCORE.equals(order.getProperty())) {
                    ScoreSortBuilder sort = SortBuilders //
                            .scoreSort() //
                            .order(sortOrder);

                    searchRequestBuilder.addSort(sort);
                } else {
                    FieldSortBuilder sort = SortBuilders //
                            .fieldSort(order.getProperty()) //
                            .order(sortOrder);

                    if (order.getNullHandling() == Sort.NullHandling.NULLS_FIRST) {
                        sort.missing("_first");
                    } else if (order.getNullHandling() == Sort.NullHandling.NULLS_LAST) {
                        sort.missing("_last");
                    }

                    searchRequestBuilder.addSort(sort);
                }
            }
        }

        if (query.getMinScore() > 0) {
            searchRequestBuilder.setMinScore(query.getMinScore());
        }
        return searchRequestBuilder;
    }

    private static String[] toArray(List<String> values) {
        String[] valuesAsArray = new String[values.size()];
        return values.toArray(valuesAsArray);
    }
}
